package com.gobang.common.interceptor;

import com.gobang.service.AuthService;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageChannel;
import org.springframework.messaging.simp.stomp.StompCommand;
import org.springframework.messaging.simp.stomp.StompHeaderAccessor;
import org.springframework.messaging.support.ChannelInterceptor;
import org.springframework.messaging.support.MessageHeaderAccessor;
import org.springframework.stereotype.Component;

/**
 * WebSocket认证拦截器
 * 拦截WebSocket STOMP消息并验证JWT Token
 * 
 * @author System
 * @since 1.0.0
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class WebSocketAuthInterceptor implements ChannelInterceptor {

    private final AuthService authService;

    /**
     * 在消息（信道消息）发送到通道前拦截处理
     *
     * @param message
     * @param channel
     * @return
     */
    @Override
    public Message<?> preSend(Message<?> message, MessageChannel channel) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
        
        if (accessor != null) {
            StompCommand command = accessor.getCommand();
            
            // 只对CONNECT和SEND命令进行认证
            if (StompCommand.CONNECT.equals(command) || StompCommand.SEND.equals(command)) {
                String token = extractToken(accessor);
                
                if (token != null) {
                    // 验证Token
                    if (authService.validateToken(token)) {
                        Long userId = authService.getUserIdFromToken(token);
                        if (userId != null) {
                            // 将用户ID存储到会话属性中
                            accessor.setUser(() -> userId.toString());
                            log.debug("WebSocket认证成功，用户ID：{}，命令：{}", userId, command);
                        } else {
                            log.warn("WebSocket认证失败，无法获取用户ID，命令：{}", command);
                            return null; // 拒绝连接
                        }
                    } else {
                        log.warn("WebSocket认证失败，Token无效，命令：{}", command);
                        return null; // 拒绝连接
                    }
                } else {
                    // CONNECT命令必须提供Token
                    if (StompCommand.CONNECT.equals(command)) {
                        log.warn("WebSocket连接被拒绝，未提供Token");
                        return null; // 拒绝连接
                    }

                    // 严格意义上的：拒绝无 Token 的 send
//                    log.warn("WebSocket连接被拒绝，未提供Token");
//                    return null; // 拒绝连接
                }
            }
        }
        
        return message;
    }

    /**
     * 从STOMP头中提取JWT Token
     *
     * @param accessor STOMP头访问器
     * @return JWT Token，如果不存在返回null
     */
    private String extractToken(StompHeaderAccessor accessor) {
        // 1. 从Authorization头获取
        String authHeader = accessor.getFirstNativeHeader("Authorization");
        if (authHeader != null && authHeader.startsWith("Bearer ")) {
            return authHeader.substring(7);
        }
        
        // 2. 从token头获取
        String tokenHeader = accessor.getFirstNativeHeader("token");
        if (tokenHeader != null && !tokenHeader.trim().isEmpty()) {
            return tokenHeader;
        }
        
        // 3. 从连接参数中获取（URL参数）
        String tokenParam = accessor.getFirstNativeHeader("X-Token");
        if (tokenParam != null && !tokenParam.trim().isEmpty()) {
            return tokenParam;
        }
        
        return null;
    }

    /**
     * 后置行为：日志记录
     *
     * @param message
     * @param channel
     * @param sent
     */
    @Override
    public void postSend(Message<?> message, MessageChannel channel, boolean sent) {
        StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
        
        if (accessor != null) {
            StompCommand command = accessor.getCommand();
            
            if (StompCommand.CONNECT.equals(command)) {
                if (sent) {
                    log.debug("WebSocket连接建立成功");
                } else {
                    log.warn("WebSocket连接建立失败");
                }
            } else if (StompCommand.DISCONNECT.equals(command)) {
                log.debug("WebSocket连接断开");
            }
        }
    }

    @Override
    public boolean preReceive(MessageChannel channel) {
        return true;
    }

    @Override
    public Message<?> postReceive(Message<?> message, MessageChannel channel) {
        return message;
    }

    @Override
    public void afterSendCompletion(Message<?> message, MessageChannel channel, boolean sent, Exception ex) {
        if (ex != null) {
            StompHeaderAccessor accessor = MessageHeaderAccessor.getAccessor(message, StompHeaderAccessor.class);
            if (accessor != null) {
                StompCommand command = accessor.getCommand();
                log.error("WebSocket消息发送异常，命令：{}，错误：{}", command, ex.getMessage());
            }
        }
    }
} 